#!/bin/bash
# SSL Certificate Setup Script for ArXiv Subscription Platform

set -e

# Configuration
DOMAIN="${1:-localhost}"
EMAIL="${2:-admin@${DOMAIN}}"
CERT_DIR="/etc/ssl/certs"
KEY_DIR="/etc/ssl/private"
NGINX_DIR="/etc/nginx"
WEBROOT="/var/www/certbot"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

log() {
    echo -e "${GREEN}[$(date '+%Y-%m-%d %H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[$(date '+%Y-%m-%d %H:%M:%S')] WARNING:${NC} $1"
}

error() {
    echo -e "${RED}[$(date '+%Y-%m-%d %H:%M:%S')] ERROR:${NC} $1"
    exit 1
}

info() {
    echo -e "${BLUE}[$(date '+%Y-%m-%d %H:%M:%S')] INFO:${NC} $1"
}

check_prerequisites() {
    log "Checking SSL setup prerequisites..."
    
    # Check if running as root
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root"
    fi
    
    # Check if certbot is installed
    if ! command -v certbot &> /dev/null; then
        info "Installing certbot..."
        apt-get update
        apt-get install -y certbot python3-certbot-nginx
    fi
    
    # Check if openssl is installed
    if ! command -v openssl &> /dev/null; then
        info "Installing openssl..."
        apt-get install -y openssl
    fi
    
    # Create directories
    mkdir -p "$CERT_DIR" "$KEY_DIR" "$WEBROOT"
    chmod 755 "$CERT_DIR"
    chmod 700 "$KEY_DIR"
    
    log "Prerequisites check completed"
}

generate_dhparam() {
    log "Generating DH parameters for perfect forward secrecy..."
    
    local dhparam_file="$CERT_DIR/dhparam.pem"
    
    if [[ ! -f "$dhparam_file" ]]; then
        info "This may take several minutes..."
        openssl dhparam -out "$dhparam_file" 2048
        chmod 644 "$dhparam_file"
        log "DH parameters generated successfully"
    else
        info "DH parameters already exist, skipping generation"
    fi
}

generate_self_signed_cert() {
    log "Generating self-signed certificate for $DOMAIN..."
    
    local cert_file="$CERT_DIR/$DOMAIN.pem"
    local key_file="$KEY_DIR/$DOMAIN.key"
    
    # Generate private key
    openssl genrsa -out "$key_file" 2048
    chmod 600 "$key_file"
    
    # Generate certificate
    openssl req -new -x509 -key "$key_file" -out "$cert_file" -days 365 -subj "/C=US/ST=State/L=City/O=Organization/OU=OrgUnit/CN=$DOMAIN"
    chmod 644 "$cert_file"
    
    log "Self-signed certificate generated successfully"
}

generate_default_cert() {
    log "Generating default certificate for nginx..."
    
    local cert_file="$CERT_DIR/default.crt"
    local key_file="$KEY_DIR/default.key"
    
    if [[ ! -f "$cert_file" || ! -f "$key_file" ]]; then
        openssl req -x509 -nodes -days 365 -newkey rsa:2048 \
            -keyout "$key_file" \
            -out "$cert_file" \
            -subj "/C=US/ST=State/L=City/O=Default/CN=default"
        
        chmod 600 "$key_file"
        chmod 644 "$cert_file"
        
        log "Default certificate generated successfully"
    else
        info "Default certificate already exists"
    fi
}

setup_letsencrypt() {
    local domain="$1"
    local email="$2"
    
    log "Setting up Let's Encrypt certificate for $domain..."
    
    # Validate domain
    if [[ "$domain" == "localhost" || "$domain" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
        warn "Cannot obtain Let's Encrypt certificate for localhost or IP address"
        warn "Using self-signed certificate instead"
        generate_self_signed_cert
        return 0
    fi
    
    # Check if nginx is running
    local nginx_running=false
    if systemctl is-active --quiet nginx; then
        nginx_running=true
        info "Stopping nginx for certificate generation..."
        systemctl stop nginx
    fi
    
    # Obtain certificate
    info "Obtaining Let's Encrypt certificate..."
    if certbot certonly \
        --standalone \
        --email "$email" \
        --agree-tos \
        --no-eff-email \
        --domains "$domain,www.$domain" \
        --non-interactive \
        --staple-ocsp \
        --preferred-challenges http; then
        
        log "Let's Encrypt certificate obtained successfully"
        
        # Copy certificates to nginx directory
        cp "/etc/letsencrypt/live/$domain/fullchain.pem" "$CERT_DIR/$domain.pem"
        cp "/etc/letsencrypt/live/$domain/privkey.pem" "$KEY_DIR/$domain.key"
        
        # Set permissions
        chmod 644 "$CERT_DIR/$domain.pem"
        chmod 600 "$KEY_DIR/$domain.key"
        
        # Setup auto-renewal
        setup_auto_renewal "$domain"
        
    else
        error "Failed to obtain Let's Encrypt certificate"
    fi
    
    # Restart nginx if it was running
    if [[ "$nginx_running" == true ]]; then
        info "Starting nginx..."
        systemctl start nginx
    fi
}

setup_auto_renewal() {
    local domain="$1"
    
    log "Setting up automatic certificate renewal..."
    
    # Create renewal script
    cat > /usr/local/bin/renew-cert.sh << EOF
#!/bin/bash
# Auto-renewal script for Let's Encrypt certificates

set -e

DOMAIN="$domain"
NGINX_RUNNING=false

# Check if nginx is running
if systemctl is-active --quiet nginx; then
    NGINX_RUNNING=true
fi

# Stop nginx if running
if [[ "\$NGINX_RUNNING" == true ]]; then
    systemctl stop nginx
fi

# Renew certificate
if certbot renew --quiet --force-renewal --domains "\$DOMAIN,www.\$DOMAIN"; then
    echo "Certificate renewed successfully"
    
    # Copy new certificates
    cp "/etc/letsencrypt/live/\$DOMAIN/fullchain.pem" "$CERT_DIR/\$DOMAIN.pem"
    cp "/etc/letsencrypt/live/\$DOMAIN/privkey.pem" "$KEY_DIR/\$DOMAIN.key"
    
    # Set permissions
    chmod 644 "$CERT_DIR/\$DOMAIN.pem"
    chmod 600 "$KEY_DIR/\$DOMAIN.key"
    
    # Reload nginx if it was running
    if [[ "\$NGINX_RUNNING" == true ]]; then
        systemctl start nginx
        systemctl reload nginx
    fi
    
    echo "Certificate renewal completed"
else
    echo "Certificate renewal failed"
    exit 1
fi
EOF

    chmod +x /usr/local/bin/renew-cert.sh
    
    # Create systemd timer for auto-renewal
    cat > /etc/systemd/system/cert-renewal.service << EOF
[Unit]
Description=Renew SSL certificates
After=network.target

[Service]
Type=oneshot
ExecStart=/usr/local/bin/renew-cert.sh
User=root
EOF

    cat > /etc/systemd/system/cert-renewal.timer << EOF
[Unit]
Description=Renew SSL certificates daily
Requires=cert-renewal.service

[Timer]
OnCalendar=daily
RandomizedDelaySec=3600
Persistent=true

[Install]
WantedBy=timers.target
EOF

    # Enable and start timer
    systemctl daemon-reload
    systemctl enable cert-renewal.timer
    systemctl start cert-renewal.timer
    
    log "Auto-renewal setup completed"
}

test_certificate() {
    local domain="$1"
    
    log "Testing SSL certificate for $domain..."
    
    local cert_file="$CERT_DIR/$domain.pem"
    local key_file="$KEY_DIR/$domain.key"
    
    if [[ -f "$cert_file" && -f "$key_file" ]]; then
        # Check certificate validity
        if openssl x509 -in "$cert_file" -text -noout > /dev/null 2>&1; then
            info "Certificate is valid"
            
            # Show certificate details
            info "Certificate details:"
            openssl x509 -in "$cert_file" -text -noout | grep -E "(Subject:|Issuer:|Not Before:|Not After :)"
            
            # Test certificate and key match
            cert_hash=$(openssl x509 -noout -modulus -in "$cert_file" | openssl md5)
            key_hash=$(openssl rsa -noout -modulus -in "$key_file" | openssl md5)
            
            if [[ "$cert_hash" == "$key_hash" ]]; then
                info "Certificate and key match"
            else
                error "Certificate and key do not match"
            fi
            
        else
            error "Certificate is invalid"
        fi
    else
        error "Certificate files not found"
    fi
}

update_nginx_config() {
    local domain="$1"
    
    log "Updating nginx configuration for SSL..."
    
    # Backup existing configuration
    if [[ -f "$NGINX_DIR/nginx.conf" ]]; then
        cp "$NGINX_DIR/nginx.conf" "$NGINX_DIR/nginx.conf.backup.$(date +%Y%m%d_%H%M%S)"
    fi
    
    # Update SSL configuration in nginx
    sed -i "s/your-domain\.com/$domain/g" "$NGINX_DIR/nginx-ssl.conf"
    
    # Test nginx configuration
    if nginx -t; then
        info "Nginx configuration is valid"
    else
        error "Nginx configuration test failed"
    fi
}

create_monitoring_script() {
    log "Creating SSL certificate monitoring script..."
    
    cat > /usr/local/bin/check-ssl-expiry.sh << 'EOF'
#!/bin/bash
# SSL Certificate Expiry Monitoring Script

set -e

DOMAIN="${1:-localhost}"
CERT_FILE="/etc/ssl/certs/$DOMAIN.pem"
DAYS_WARNING=30
EMAIL_ALERT="${EMAIL_ALERT:-admin@example.com}"

if [[ ! -f "$CERT_FILE" ]]; then
    echo "Certificate file not found: $CERT_FILE"
    exit 1
fi

# Check certificate expiry
EXPIRY_DATE=$(openssl x509 -in "$CERT_FILE" -noout -dates | grep "notAfter" | cut -d= -f2)
EXPIRY_EPOCH=$(date -d "$EXPIRY_DATE" +%s)
CURRENT_EPOCH=$(date +%s)
DAYS_UNTIL_EXPIRY=$(( (EXPIRY_EPOCH - CURRENT_EPOCH) / 86400 ))

echo "Certificate expires in $DAYS_UNTIL_EXPIRY days ($EXPIRY_DATE)"

if [[ $DAYS_UNTIL_EXPIRY -le $DAYS_WARNING ]]; then
    echo "WARNING: Certificate expires in $DAYS_UNTIL_EXPIRY days!"
    
    # Send alert email if configured
    if command -v mail &> /dev/null && [[ -n "$EMAIL_ALERT" ]]; then
        echo "SSL certificate for $DOMAIN expires in $DAYS_UNTIL_EXPIRY days. Please renew it." | \
            mail -s "SSL Certificate Expiry Warning - $DOMAIN" "$EMAIL_ALERT"
    fi
    
    # Log to syslog
    logger -t ssl-monitor "SSL certificate for $DOMAIN expires in $DAYS_UNTIL_EXPIRY days"
    
    exit 1
fi

echo "Certificate is valid"
EOF

    chmod +x /usr/local/bin/check-ssl-expiry.sh
    
    # Add to cron for daily monitoring
    (crontab -l 2>/dev/null; echo "0 6 * * * /usr/local/bin/check-ssl-expiry.sh $DOMAIN") | crontab -
    
    log "SSL monitoring setup completed"
}

show_summary() {
    log "SSL Setup Summary"
    echo
    info "Domain: $DOMAIN"
    info "Certificate: $CERT_DIR/$DOMAIN.pem"
    info "Private Key: $KEY_DIR/$DOMAIN.key"
    info "DH Parameters: $CERT_DIR/dhparam.pem"
    echo
    info "Auto-renewal: Enabled (daily check)"
    info "Monitoring: Enabled (daily expiry check)"
    echo
    log "SSL setup completed successfully!"
    echo
    info "Next steps:"
    info "1. Update your nginx configuration to use SSL"
    info "2. Test your website with: curl -I https://$DOMAIN"
    info "3. Check SSL rating: https://www.ssllabs.com/ssltest/"
    echo
}

# Main execution
main() {
    if [[ -z "$DOMAIN" || "$DOMAIN" == "localhost" ]]; then
        warn "No valid domain provided. Using localhost with self-signed certificate."
        DOMAIN="localhost"
    fi
    
    log "Starting SSL setup for $DOMAIN..."
    
    check_prerequisites
    generate_dhparam
    generate_default_cert
    
    if [[ "$DOMAIN" != "localhost" ]]; then
        setup_letsencrypt "$DOMAIN" "$EMAIL"
    else
        generate_self_signed_cert
    fi
    
    test_certificate "$DOMAIN"
    create_monitoring_script
    show_summary
}

# Script execution
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi